package org.emailclassifier.application;

import java.util.Scanner;
import org.emailclassifier.classifiers.AIMClassifier;
//import org.emailclassifier.classifiers.ClassProbComparator;
import org.emailclassifier.classifiers.ClassifierBase;
import org.emailclassifier.classifiers.NBClassifier;
import org.emailclassifier.classifiers.SVMClassifier;
import org.emailclassifier.extractdb.BatchProcessor;
import org.emailclassifier.extractdb.EmailParser;
import org.emailclassifier.preprocessing.DataSetDocument;
import org.emailclassifier.preprocessing.EmailDataSet;

/**
 * Entry point of the application
 */
public class Main {
	
	/**
	 * Entry point of the application
	 * @param args Command-line arguments
	 */
	public static void main( String[] args ) {
		
		if ( args.length != 1 ) {
			System.out.println( "Missing parameter: dataset directory" );
		}
		else {
			showMenu( args[0] );
		}
		
	}
	
    private static int showMenu( String emaildb )
    {
 
	    Scanner scanner = new Scanner (System.in);
	    int choice = 0;
	    boolean choiceOk = false;
	    double cvf = 0;
	    int iters = 0;
	    int dataLevel = 0;
	    ClassifierBase classifier = null;
	    	    
		System.out.println("ENRON E-MAIL CLASSIFIER");
		System.out.println("By:  Lourens Elzinga");
		System.out.println("     Ante Znaor");
		System.out.println("***********************"); 
	    
		System.out.println("\n\n+---------------------+"); 
		System.out.println("| Select a classifier |"); 
		System.out.println("+---------------------+\n");  
	    System.out.println("[1] Naive Bayes Classifier"); 
	    System.out.println("[2] AIM Classifier");
	    System.out.println("[3] SVM Classifier");
	    System.out.println("[4] Exit"); 
	    
	    while(choiceOk == false) {
	        System.out.print("Selection: ");
	        choice = scanner.nextInt();
	        System.out.println("");
	        
	        switch (choice) {
        	case 1: 
        		classifier = new NBClassifier();
        		choiceOk = true;
        	   	break;
	        
        	case 2:
        		classifier = new AIMClassifier();
        		choiceOk = true;
	            break;	         
	    
        	case 3:
        		classifier = new SVMClassifier();
        		choiceOk = true;
	            break;	 	            
	            
	        case 4:
	        	System.out.println("Exiting...");
	            System.exit(0);
	            break;
	                    
	        default:
	        	System.out.println("Please enter a valid selection!");
	        	break;
		    }
	    }
	    
	    choiceOk = false;
	     
		System.out.println("Choose data to train on:"); 
	    System.out.println("[1] Email body"); 
	    System.out.println("[2] Email body, sender and subject"); 
	    System.out.println("[3] Email body and all headers");
	    while(choiceOk == false) {
	        System.out.print("Selection: ");
	        choice = scanner.nextInt();
	        System.out.println("");
	        
	        switch (choice)
	        {
        	case 1: 
        	case 2:    
	        case 3:
	        	dataLevel = choice - 1;
	        	choiceOk = true;
	            break;
	                    
	        default:
	        	System.out.println("Please enter a valid selection!");
	        	break;
		    }
	    }
	    
	    System.out.print("Use stoplist (Y/N): ");
	    boolean stoplist = scanner.next().equalsIgnoreCase("y");
	    
	    	    
	    System.out.print("Enter the cross-validation fraction (0-100): ");
		cvf = new Double(scanner.nextInt()) / 100;
		if (cvf < 0) cvf = 0;
		if (cvf > 1) cvf = 1;
		
		System.out.print("Enter the amount of iterations (min. 1): ");
		iters = new Integer(scanner.nextInt());
		if (iters < 1) iters = 1;
	    
		testClassifier( classifier, emaildb, cvf, iters, dataLevel, stoplist);
		
	    return 0;
    }
	

	/**
	 * Tests the given classifier
	 * @param classifier The classifier to test
	 * @param dataSetRoot Root directory of the data set
	 * @param crossValidationFraction The fraction of the data set to use as training set (0.0 to 1.0)
	 */
	static public void testClassifier( ClassifierBase classifier, String dataSetRoot, double crossValidationFraction, int iterations, int dataLevel, boolean stoplist )
	{
		// Benchmarking variables
		final long loadStartTime;
		final long loadEndTime;
		
		System.out.printf("Testing classifier: %s\n", classifier.getDescription());
		
		
		System.out.printf("Loading training samples from %s\n", dataSetRoot);
		
		loadStartTime = System.nanoTime();
		
		BatchProcessor enronDataSet = new BatchProcessor();
		enronDataSet.process(dataSetRoot);
		
		String[] categories = enronDataSet.getCategories();
		
		EmailDataSet dataSet = new EmailDataSet();
		dataSet.setStoplist(stoplist);
		
		for( String cat : categories ) {
			EmailParser[] enronMails = enronDataSet.getEmails( cat ); 
			
			for ( EmailParser enronMail : enronMails ) {
				switch (dataLevel)
				{
				case 0:
					dataSet.addDocument(cat, enronMail.getBody());
					break;
				case 1:
					dataSet.addDocument(cat, enronMail.getBody() + " " + enronMail.getSender() + " " + enronMail.getSubject());
					break;
				case 2:
					dataSet.addDocument(cat, enronMail.getRawMail());
					break;
				}				
			}
		}
		
		loadEndTime = System.nanoTime();
		
		System.out.printf("Loaded dataset with %d samples in %d categorories\n", dataSet.getGlobalDocumentCount(), dataSet.getCategoryCount());
		
		int correct = 0;
		int total = 0;
		int[] correctRank = new int[categories.length];
		double SEMData[][] = new double[iterations][categories.length];
		double SEM[] = new double[categories.length];
				
		long crossTime = 0;
		long trainTime = 0;
		long testTime = 0;
		for (int iteration = 0; iteration < iterations; iteration++) {
			
			// Benchmarking variables
			final long crossStartTime;
			final long crossEndTime;
			
			final long trainStartTime;
			final long trainEndTime;
			
			final long testStartTime;
			final long testEndTime;
			
			int[] correctRankIteration = new int[categories.length]; 
			int totalIteration = 0;
		
			System.out.printf("Running test iteration %d of %d\n", iteration + 1, iterations);
			
			System.out.printf("Generating training and testing set with cross-validation (%f fraction)...\n", crossValidationFraction );
			
			crossStartTime = System.nanoTime();
			EmailDataSet[] cv = dataSet.makeCrossValidationSet(crossValidationFraction);
			
			EmailDataSet trainSet = cv[0];
			EmailDataSet testSet = cv[1];
			
			crossEndTime = System.nanoTime();
			crossTime += (crossEndTime - crossStartTime);
			
			classifier.setTrainingSet(trainSet);
			
			System.out.printf("Training %s\n", classifier.getDescription());
			
			trainStartTime = System.nanoTime();
			classifier.train();
			trainEndTime = System.nanoTime();
			
			trainTime += (trainEndTime - trainStartTime);
			
			System.out.printf("Testing %s\n", classifier.getDescription());
						
			testStartTime = System.nanoTime();
			
			for( String cat : categories ) {
				
				for ( DataSetDocument doc : testSet.getDocumentsInCategory(cat) ) {
					
					String results[] = classifier.determineClasses( doc );
					for (int i = 0; i < categories.length; i++)
					{
						if (results[i] == cat) {
							correctRank[i]++;
							correctRankIteration[i]++;
							if (i == 0) {
								correct++;
							}
						}
					}
					total++;
					totalIteration++;
				}
			}
			
			// Update SEM data
			for (int c = 0; c < categories.length; c++ )
			{
				SEMData[iteration][c] = (double)correctRankIteration[c] / (double)totalIteration * 100;
			}
			
			testEndTime = System.nanoTime();
			testTime += (testEndTime - testStartTime);
		}
		
		
		for (int c = 0; c < categories.length; c++)
		{
			double catMean = (double)correctRank[c]/(double)total*100;
			double MSE = 0;
			for (int i = 0; i < iterations; i++)
			{
				MSE += Math.pow(SEMData[i][c] - catMean, 2);
			}
			SEM[c] = Math.sqrt( (1 / ((double)iterations - 1)) * MSE ) / Math.sqrt(iterations);
		}		
		
		System.out.print("Testing finished, outputting results:\n");
		
		double cummulative = 0;
		int cumQty = 0;
		System.out.printf("Rank\tCorrect (%%)\tSEM\tCummulative (%%)\tQty\tCumQty\n");
		for (int i = 0; i<categories.length; i++)
		{
			cummulative += (double)correctRank[i]/(double)total*100;
			cumQty += correctRank[i];
			
			System.out.printf("%d\t%f%%\t%f%%\t%f%%\t%d\t%d\n", i+1, (double)correctRank[i]/(double)total*100, SEM[i], cummulative, correctRank[i], cumQty);
		}
		
		
		System.out.printf("%d out of %d tests correct (%f%% accuracy)\n", correct, total, (double)correct/(double)total*100);
		
		// Output performance information
		System.out.printf("Performance:\n");
		final double loadTime = (double)(loadEndTime - loadStartTime) / 1000000;
		System.out.printf("\tLoading dataset: %f ms\n", loadTime );
		final double cvTime = (double)(crossTime) / (1000000 * iterations);
		System.out.printf("\tCreating cross-validation set: %f ms\n", cvTime );
		final double trTime = (double)(trainTime) / (1000000 * iterations);
		System.out.printf("\tTraining: %f ms\n", trTime );
		final double tsTime = (double)(testTime) / (1000000 * iterations);
		System.out.printf("\tTesting: %f ms\n", tsTime );
	}
}